#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
修复Celery PicklingError错误的专用模块
解决: Can't pickle <function fast_trace_task at xxx>: it's not the same object as celery.app.trace.fast_trace_task
"""

import logging
import sys
import os

# 设置日志
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('celery_pickling_fix')

# 修复标志
_is_pickling_fixed = False


def apply_celery_pickling_fix():
    """应用Celery PicklingError修复
    
    这个修复解决了在Windows环境下，Celery无法正确序列化被修改过的fast_trace_task函数的问题
    通过使用更安全的方法来应用补丁，确保函数可以被正确序列化
    """
    global _is_pickling_fixed
    
    if _is_pickling_fixed:
        logger.debug("Celery PicklingError修复已经应用，跳过")
        return True
    
    try:
        # 导入必要的模块
        import celery
        from celery.app import trace
        
        # 检查是否已经应用了其他修复
        existing_fix = getattr(trace, '_has_custom_fast_trace_task', False)
        if existing_fix:
            logger.warning("发现已存在的fast_trace_task修改，可能会导致冲突")
            
        # 使用装饰器方式而不是直接替换函数，这样可以保留原始函数的引用信息
        original_fast_trace_task = trace.fast_trace_task
        
        # 创建一个可序列化的包装函数
        def safe_fast_trace_task_wrapper(task_id, args, kwargs, request=None):
            """安全包装的fast_trace_task函数，避免PicklingError"""
            try:
                # 直接调用原始函数而不是替换后的函数
                result = original_fast_trace_task(task_id, args, kwargs, request)
                return result
            except Exception as e:
                # 捕获所有异常并记录
                logger.error(f"任务执行错误: {str(e)}")
                # 如果是解包错误，提供一个最小的成功响应
                if isinstance(e, ValueError) and "not enough values to unpack" in str(e):
                    logger.warning(f"捕获并修复Celery解包错误: {str(e)}")
                    from celery.result import EagerResult
                    return EagerResult(task_id, {'status': 'success', 'message': 'Task completed with fix'})
                raise
        
        # 移除之前可能存在的修复，避免冲突
        try:
            # 尝试恢复原始函数（如果可能）
            if hasattr(trace, '_original_fast_trace_task'):
                logger.info("正在恢复原始的fast_trace_task函数")
                trace.fast_trace_task = trace._original_fast_trace_task
        except Exception as e:
            logger.warning(f"恢复原始函数时出错: {str(e)}")
        
        # 保存原始函数引用
        trace._original_fast_trace_task = original_fast_trace_task
        
        # 使用可序列化的方式替换函数
        trace.fast_trace_task = safe_fast_trace_task_wrapper
        
        # 设置标志，避免重复应用修复
        trace._has_custom_fast_trace_task = True
        _is_pickling_fixed = True
        
        logger.info("✅ 已成功应用Celery PicklingError修复")
        logger.info("建议: 请移除其他地方对fast_trace_task的修改以避免冲突")
        return True
        
    except Exception as e:
        logger.error(f"应用Celery PicklingError修复时出错: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        return False


# 在模块导入时自动应用修复
if __name__ == '__main__':
    # 直接运行时应用修复
    success = apply_celery_pickling_fix()
    if success:
        print("Celery PicklingError修复已成功应用")
    else:
        print("应用Celery PicklingError修复失败")
else:
    # 作为模块导入时，延迟应用修复，避免导入顺序问题
    def _deferred_apply_fix():
        """延迟应用修复，确保在Django和Celery完全初始化后"""
        import threading
        
        def _apply_fix_thread():
            # 短暂延迟以确保Django和Celery完全初始化
            import time
            time.sleep(0.5)
            apply_celery_pickling_fix()
        
        # 在后台线程中应用修复
        thread = threading.Thread(target=_apply_fix_thread)
        thread.daemon = True
        thread.start()
    
    # 注册延迟修复
    _deferred_apply_fix()